#Import libraries and setup matplotlib
import matplotlib
# %matplotlib inline
import matplotlib.pylab as plt
import librosa
import IPython.display as ipd

import sys
sys.path.append('waveglow/')
import numpy as np
import torch

# from hparams import create_hparams
from hparams_tts import create_hparams
from model import Tacotron2
from layers import TacotronSTFT, STFT
from audio_processing import griffin_lim
# from train import load_model
# from text import text_to_sequence
from denoiser import Denoiser
from train_tts import load_model
from text import text_to_sequence
from waveglow.denoiser import Denoiser

def plot_data(data, figsize=(16, 4)):
    fig, axes = plt.subplots(1, len(data), figsize=figsize)
    for i in range(len(data)):
        axes[i].imshow(data[i], aspect='auto', origin='bottom', 
                       interpolation='none')
# Setup hparams
hparams = create_hparams()
hparams.sampling_rate = 22050
# Load model from checkpoint
checkpoint_path = "models/tacotron2_statedict.pt"
model = load_model(hparams)
model.load_state_dict(torch.load(checkpoint_path)['state_dict'])
_ = model.cuda().eval().half()
# Load WaveGlow for mel2audio synthesis and denoiser
waveglow_path = 'models/waveglow_256channels_universal_v5.pt'
waveglow = torch.load(waveglow_path)['model']
waveglow.cuda().eval().half()
for k in waveglow.convinv:
    k.float()
denoiser = Denoiser(waveglow)
# Prepare text input
# text = "Waveglow is really awesome!"
text = "Hello Musixmatch, how are you?"
sequence = np.array(text_to_sequence(text, ['english_cleaners']))[None, :]
sequence = torch.autograd.Variable(torch.from_numpy(sequence)).cuda().long()

# Decode text input and plot results
mel_outputs, mel_outputs_postnet, _, alignments = model.inference(sequence)

# mel_outputs = mel_outputs.float().data.cpu().numpy()[0]
# mel_outputs_postnet = mel_outputs_postnet.float().data.cpu().numpy()[0]
# alignments= alignments.float().data.cpu().numpy()[0].T


# plot_data((mel_outputs.float().data.cpu().numpy()[0],
#            mel_outputs_postnet.float().data.cpu().numpy()[0],
#            alignments.float().data.cpu().numpy()[0].T))


# Synthesize audio from spectrogram using WaveGlow

with torch.no_grad():
    audio = waveglow.infer(mel_outputs_postnet, sigma=0.666)

ipd.Audio(audio[0].data.cpu().numpy(), rate=hparams.sampling_rate)
# (Optional) Remove WaveGlow bias
audio_denoised = denoiser(audio, strength=0.01)[:, 0]
ipd.Audio(audio_denoised.cpu().numpy(), rate=hparams.sampling_rate)
# save
audio = audio.cpu().numpy()
audio = audio.astype('float64')
outAudioPath = "test_1.wav"
librosa.output.write_wav(outAudioPath, audio[0], hparams.sampling_rate, norm=False)